# 导入大模型工具
from llm_tool import LLMTool
# 导入工具类
from utils import Utils
# 导入matlab工具类
from matlab_tool import MatlabTool
import pandas as pd
import shutil
import glob
# 设置实验文件路径
requirement_file_path = "./dataset/lm_challenges/original_models/0_triplex/0_triplex_reqs.docx"
# 设置当前实验编号
experiment_id = "910-2"
# 初始化大模型上下文
llm_tool = LLMTool()
def main():
    
    
   # 设置实验的模型路径,并将模型文件复制到实验文件夹
    model_path = "./dataset/lm_challenges/original_models/0_triplex/triplex_12B.mdl"
    Utils.copy_files(model_path, f"./experiments/{experiment_id}/model")
    script_path = "./dataset/lm_challenges/original_models/0_triplex/run.m"
    Utils.copy_files(script_path, f"./experiments/{experiment_id}/model")
    # 1、与大模型交互，识别需求领域并抽取需求条目
    interact_with_llm(prompt_path="./prompts/s1_base_info.txt", 
                      ai_response_path=f"./experiments/{experiment_id}/ai_response_s1.txt",
                        step="step 1")

    # 2、与大模型交互, 识别该系统的状态并分析迁移条件
    interact_with_llm(prompt_path="./prompts/s2_scenario_config.txt", 
                      ai_response_path= f"./experiments/{experiment_id}/ai_response_s2.txt",
                        step="step 2")

    # 3、与大模型交互，检查生成的状态是否合理且是否有遗漏情况，如果有，更新2中的输出
    interact_with_llm(prompt_path="./prompts/s3_state_check.txt", 
                      ai_response_path= f"./experiments/{experiment_id}/ai_response_s3.txt",
                        step="step 3")
    # 4、与大模型交互，结合需求条目生成场景
    interact_with_llm(prompt_path="./prompts/s4_scenario_generation.txt", 
                      ai_response_path= f"./experiments/{experiment_id}/ai_response_s4.txt",
                        step="step 4")

    # 5、与大模型交互，生成测试用例
    interact_with_llm(prompt_path="./prompts/s5_test_case_generation.txt", 
                      ai_response_path= f"./experiments/{experiment_id}/ai_response_s5.txt",
                        step="step 5")


    # 6、与大模型交互，整理用例为csv文本
    interact_with_llm(prompt_path="./prompts/s6_test_case2csv.txt", 
                      ai_response_path= f"./experiments/{experiment_id}/ai_response_s6.txt",
                        step="step 6")
    # 6.1、提取ai_response_s5中的csv代码块，将其保存为.csv文件
    csv_file_path = f"./experiments/{experiment_id}/testcase/testcase_0.csv"
    Utils.extract_csv_from_response(f"./experiments/{experiment_id}/ai_response_s6.txt", csv_file_path)
    # 将testcase复制到模型目录
    shutil.copy(csv_file_path, f"./experiments/{experiment_id}/model/testcase.csv")

    # 7、启动matlab引擎，运行simulink模型, 获取仿真结果
    matlab_tool = MatlabTool()
    script_path = "run"
    script_log_path = f"./experiments/{experiment_id}/model/run.log"
    model_dir = f"./experiments/{experiment_id}/model"
    # 后续记得更改modelname
    modelName = "'triplex_12B'"
    csv_file_path = f"./experiments/{experiment_id}/model/testcase.csv"
    # 获取仿真周期数
    df = pd.read_csv(csv_file_path)
    simulation_time = len(df)
    # 返回仿真结果，保存到字符串中
    matlab_tool.run_mscript(script_path, script_log_path, model_dir, modelName, simulation_time)
    simulation_results = matlab_tool.get_simulation_results(model_dir)

    # 反馈优化机制
    reward_times = 3
    # 7、反馈大模型仿真结果，同时重新运行模型
    for i in range(reward_times):
        # 7.1 反馈大模型仿真结果重新生成测试场景
        interact_with_llm(prompt_path="./prompts/s7_feedback_scenario_generation.txt",
                          ai_response_path= f"./experiments/{experiment_id}/ai_response_reward{i+1}_s7-1.txt",
                          step=f"reward{i+1}: step 7-1", 
                          add_prompt=simulation_results)
        # 7.2 根据场景生成新的测试用例
        interact_with_llm(prompt_path="./prompts/s7_feedback_test_case_generation.txt",
                          ai_response_path= f"./experiments/{experiment_id}/ai_response_reward{i+1}_s7-2.txt",
                          step=f"reward{i+1}: step 7-2")
        # 7.3 将测试用例转换成csv文件
        interact_with_llm(prompt_path="./prompts/s6_test_case2csv.txt",
                          ai_response_path= f"./experiments/{experiment_id}/ai_response_reward{i+1}_s7-3.txt",
                          step=f"reward{i+1}: step 7-3")
        # 将仿真结果转换成csv格式
        csv_file_path = f"./experiments/{experiment_id}/testcase/testcase_{i+1}.csv"
        Utils.extract_csv_from_response(f"./experiments/{experiment_id}/ai_response_reward{i+1}_s7-3.txt", csv_file_path)
        shutil.copy(csv_file_path, f"./experiments/{experiment_id}/model/testcase.csv")
        print(f"reward{i+1}: csv格式转换成功")
        # 获取仿真周期数，并更新time周期数
        df = pd.read_csv(csv_file_path)
        df.iloc[:, 0] = range(1, len(df) + 1)
        df.to_csv(f"./experiments/{experiment_id}/model/testcase.csv", index=False)
        simulation_time = len(df)
        # 7.4 重新运行模型
        matlab_tool.run_mscript(script_path, script_log_path, model_dir, modelName, simulation_time)
        simulation_results = matlab_tool.get_simulation_results(model_dir)
        print(f"reward{i+1}: 重新运行模型成功")

    #合并testcase文件，获取全局的覆盖率
    print("合并所有的testcase文件，重新运行模型")
    all_testcase_path = f"./experiments/{experiment_id}/testcase/testcase_all.csv"
    files = glob.glob(f"./experiments/{experiment_id}/testcase/*.csv")
    df = pd.concat([pd.read_csv(f) for f in files], ignore_index=True)
    df.iloc[:, 0] = range(1, len(df) + 1)
    df.to_csv(all_testcase_path, index=False)
    simulation_time = len(df)
    shutil.copy(all_testcase_path, f"./experiments/{experiment_id}/model/testcase.csv")
    matlab_tool.run_mscript(script_path, script_log_path, model_dir, modelName, simulation_time)


    # 关闭matlab引擎
    matlab_tool.eng.quit()
    # 将当前prompts文件夹复制放到实验文件夹下
    Utils.copy_files("./prompts", f"./experiments/{experiment_id}/prompts")
    print("All steps done!")



# 与大模型交互的函数
def interact_with_llm(prompt_path="", ai_response_path="", step="",add_prompt=""):
    with open(prompt_path, "r", encoding="utf-8") as f:
        prompt = f.read()
    if (add_prompt != ""):
        prompt = prompt + add_prompt
    print(f"Starting {step}...")
    llm_tool.chat_with_file(ai_response_path, file_path=requirement_file_path, user_message=prompt )
    print(f"{step} done")

    

if __name__ == "__main__":
    main()
    # matlab_tool = MatlabTool()
    # model_dir = f"./experiments/{experiment_id}/model"
    # simulation_results = matlab_tool.get_simulation_results(model_dir)
